/**

  Virtual memory functions.

  by dmazar

**/

#include <Library/UefiLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/DebugLib.h>

#include "VMem.h"
#include "Lib.h"

/** Memory allocation for VM map pages that we will create with VmMapVirtualPage.
  * We need to have it preallocated during boot services.
  */
UINT8   *VmMemoryPool = NULL;
INTN    VmMemoryPoolFreePages = 0;

VOID
GetCurrentPageTable (
  PAGE_MAP_AND_DIRECTORY_POINTER  **PageTable,
  UINTN                           *Flags
) {
  UINTN CR3;

  CR3 = AsmReadCr3 ();
  //DBG ("GetCurrentPageTable: CR3 = 0x%lx\n", CR3);
  *PageTable = (PAGE_MAP_AND_DIRECTORY_POINTER *)(UINTN)(CR3 & CR3_ADDR_MASK);
  *Flags = CR3 & (CR3_FLAG_PWT | CR3_FLAG_PCD);
}

VOID
PrintPageTablePTE (
  PAGE_TABLE_4K_ENTRY   *PTE,
  VIRTUAL_ADDR          VA
) {
#if DBG_TO
  UINTN             Index;
  UINT64            Start;

  for (Index = 0; Index < 10; Index++) {
    VA.Pg4K.PTOffset = Index;
    //DBG ("      PTE %03x at %p = %lx => ", Index, PTE, PTE->Uint64);
    // 4KB PTE
    Start = (PTE->Uint64 & PT_ADDR_MASK_4K);
    //DBG ("4KB Fl: %lx VA: %lx - %lx ==> PH %lx - %lx\n",
    //  (PTE->Uint64 & ~PT_ADDR_MASK_4K), VA.Uint64, VA.Uint64 + 0x1000 - 1, Start, Start + 0x1000 - 1);
    PTE++;
  }
  //WaitForKeyPress (L"more ...");
#endif
}

VOID
PrintPageTablePDE (
  PAGE_MAP_AND_DIRECTORY_POINTER  *PDE,
  VIRTUAL_ADDR                    VA
) {
#if DBG_TO
  UINTN                 Index;
  PAGE_TABLE_2M_ENTRY   *PT2M;
  PAGE_TABLE_4K_ENTRY   *PTE;
  UINT64                Start;

  for (Index = 0; Index < 10; Index++) {
    VA.Pg4K.PDOffset = Index;
    //DBG ("    PDE %03x at %p = %lx => ", Index, PDE, PDE->Uint64);
    if (PDE->Bits.MustBeZero & 0x1) {
      // 2MB PDE
      PT2M = (PAGE_TABLE_2M_ENTRY *)PDE;
      Start = (PT2M->Uint64 & PT_ADDR_MASK_2M);
      //DBG ("2MB Fl: %lx VA: %lx - %lx ==> PH %lx - %lx\n",
      //  (PT2M->Uint64 & ~PT_ADDR_MASK_2M), VA.Uint64, VA.Uint64 + 0x200000 - 1, Start, Start + 0x200000 - 1);
    } else {
      //DBG ("  Fl: %lx %lx ->\n", (PDE->Uint64 & ~PT_ADDR_MASK_4K), (PDE->Uint64 & PT_ADDR_MASK_4K));
      PTE = (PAGE_TABLE_4K_ENTRY *)(PDE->Uint64 & PT_ADDR_MASK_4K);
      PrintPageTablePTE (PTE, VA);
    }

    PDE++;
  }
  //WaitForKeyPress (L"more ...");
#endif
}

VOID
PrintPageTablePDPE (
  PAGE_MAP_AND_DIRECTORY_POINTER  *PDPE,
  VIRTUAL_ADDR                    VA
) {
#if DBG_TO
  UINTN                           Index;
  PAGE_TABLE_1G_ENTRY             *PT1G;
  PAGE_MAP_AND_DIRECTORY_POINTER  *PDE;
  UINT64                          Start;

  for (Index = 0; Index < 10; Index++) {
    VA.Pg4K.PDPOffset = Index;
    //DBG ("  PDPE %03x at %p = %lx => ", Index, PDPE, PDPE->Uint64);
    if (PDPE->Bits.MustBeZero & 0x1) {
      // 1GB PDPE
      PT1G = (PAGE_TABLE_1G_ENTRY *)PDPE;
      Start = (PT1G->Uint64 & PT_ADDR_MASK_1G);
      //DBG ("1GB Fl: %lx VA: %lx - %lx ==> PH %lx - %lx\n",
      //  (PT1G->Uint64 & ~PT_ADDR_MASK_1G), VA.Uint64, VA.Uint64 + 0x40000000 - 1, Start, Start + 0x40000000 - 1);
    } else {
      //DBG ("  Fl: %lx %lx ->\n", (PDPE->Uint64 & ~PT_ADDR_MASK_4K), (PDPE->Uint64 & PT_ADDR_MASK_4K));
      PDE = (PAGE_MAP_AND_DIRECTORY_POINTER *)(PDPE->Uint64 & PT_ADDR_MASK_4K);
      PrintPageTablePDE (PDE, VA);
    }

    PDPE++;
  }
  //WaitForKeyPress (L"more ...");
#endif
}

VOID
PrintPageTable (
  PAGE_MAP_AND_DIRECTORY_POINTER *PageTable,
  UINTN Flags
) {
  UINTN                           Index;
  VIRTUAL_ADDR                    VA;
  PAGE_MAP_AND_DIRECTORY_POINTER  *PML4, *PDPE;

  //DBG ("PrintPageTable: %p, Flags: PWT: %d, PCD: %d\n", PageTable, (Flags & CR3_FLAG_PWT), (Flags & CR3_FLAG_PCD));
  PML4 = PageTable;

  for (Index = 0; Index < 3; Index++) {
    VA.Uint64 = 0;
    VA.Pg4K.PML4Offset = Index;
    VA_FIX_SIGN_EXTEND (VA);
    //DBG ("PML4 %03x at %p = %lx => Fl: %lx %lx ->\n", Index, PML4, PML4->Uint64, (PML4->Uint64 & ~PT_ADDR_MASK_4K), (PML4->Uint64 & PT_ADDR_MASK_4K));
    PDPE = (PAGE_MAP_AND_DIRECTORY_POINTER *)(PML4->Uint64 & PT_ADDR_MASK_4K);
    PrintPageTablePDPE (PDPE, VA);
    PML4++;
  }
  //WaitForKeyPress (L"END");
}

EFI_STATUS
GetPhysicalAddr (
  PAGE_MAP_AND_DIRECTORY_POINTER  *PageTable,
  EFI_VIRTUAL_ADDRESS             VirtualAddr,
  EFI_PHYSICAL_ADDRESS            *PhysicalAddr
) {
  EFI_PHYSICAL_ADDRESS            Start;
  VIRTUAL_ADDR                    VA, VAStart, VAEnd;
  PAGE_MAP_AND_DIRECTORY_POINTER  *PML4, *PDPE, *PDE;
  PAGE_TABLE_4K_ENTRY             *PTE4K;
  PAGE_TABLE_2M_ENTRY             *PTE2M;
  PAGE_TABLE_1G_ENTRY             *PTE1G;

  VA.Uint64 = (UINT64)VirtualAddr;
  //VA_FIX_SIGN_EXTEND (VA);
  //DBG ("PageTable: %p\n", PageTable);
  //DBG ("VA: %lx => Indexes PML4=%x, PDP=%x, PD=%x, PT=%x\n",
  //  VA.Uint64, VA.Pg4K.PML4Offset, VA.Pg4K.PDPOffset, VA.Pg4K.PDOffset, VA.Pg4K.PTOffset);

  // PML4
  PML4 = PageTable;
  PML4 += VA.Pg4K.PML4Offset;
  // prepare region start and end
  VAStart.Uint64 = 0;
  VAStart.Pg4K.PML4Offset = VA.Pg4K.PML4Offset;
  VA_FIX_SIGN_EXTEND (VAStart);
  VAEnd.Uint64 = ~(UINT64)0;
  VAEnd.Pg4K.PML4Offset = VA.Pg4K.PML4Offset;
  VA_FIX_SIGN_EXTEND (VAEnd);

  // print it
  //DBG ("PML4[%03x] at %p = %lx Region: %lx - %lx\n", VA.Pg4K.PML4Offset, PML4, PML4->Uint64, VAStart.Uint64, VAEnd.Uint64);
  if (!PML4->Bits.Present) {
    //DBG ("-> Mapping not present!\n");
    return EFI_NO_MAPPING;
  }

  //DBG ("-> Nx:%x|A:%x|PCD:%x|PWT:%x|US:%x|RW:%x|P:%x -> %lx\n",
  //PML4->Bits.Nx, PML4->Bits.Accessed,
  //PML4->Bits.CacheDisabled, PML4->Bits.WriteThrough,
  //PML4->Bits.UserSupervisor, PML4->Bits.ReadWrite, PML4->Bits.Present,
  //(PML4->Uint64 & PT_ADDR_MASK_4K)
  //);

  // PDPE
  PDPE = (PAGE_MAP_AND_DIRECTORY_POINTER *)(PML4->Uint64 & PT_ADDR_MASK_4K);
  PDPE += VA.Pg4K.PDPOffset;
  VAStart.Pg4K.PDPOffset = VA.Pg4K.PDPOffset;
  VAEnd.Pg4K.PDPOffset = VA.Pg4K.PDPOffset;
  //DBG ("PDPE[%03x] at %p = %lx Region: %lx - %lx\n", VA.Pg4K.PDPOffset, PDPE, PDPE->Uint64, VAStart.Uint64, VAEnd.Uint64);
  if (!PDPE->Bits.Present) {
    //DBG ("-> Mapping not present!\n");
    return EFI_NO_MAPPING;
  }

  if (PDPE->Bits.MustBeZero & 0x1) {
    // 1GB PDPE
    PTE1G = (PAGE_TABLE_1G_ENTRY *)PDPE;
    //DBG ("-> Nx:%x|G:%x|PAT:%x|D:%x|A:%x|PCD:%x|PWT:%x|US:%x|RW:%x|P:%x\n",
    //PTE1G->Bits.Nx, PTE1G->Bits.Global, PTE1G->Bits.PAT,
    //PTE1G->Bits.Dirty, PTE1G->Bits.Accessed,
    //PTE1G->Bits.CacheDisabled, PTE1G->Bits.WriteThrough,
    //PTE1G->Bits.UserSupervisor, PTE1G->Bits.ReadWrite, PTE1G->Bits.Present
    //);
    Start = (PTE1G->Uint64 & PT_ADDR_MASK_1G);
    *PhysicalAddr = Start + VA.Pg1G.PhysPgOffset;
    //DBG ("-> 1GB page %lx - %lx => %lx\n", Start, Start + 0x40000000 - 1, *PhysicalAddr);
    return EFI_SUCCESS;
  }
  //DBG ("-> Nx:%x|A:%x|PCD:%x|PWT:%x|US:%x|RW:%x|P:%x -> %lx\n",
  //PDPE->Bits.Nx, PDPE->Bits.Accessed,
  //PDPE->Bits.CacheDisabled, PDPE->Bits.WriteThrough,
  //PDPE->Bits.UserSupervisor, PDPE->Bits.ReadWrite, PDPE->Bits.Present,
  //(PDPE->Uint64 & PT_ADDR_MASK_4K)
  //);

  // PDE
  PDE = (PAGE_MAP_AND_DIRECTORY_POINTER *)(PDPE->Uint64 & PT_ADDR_MASK_4K);
  PDE += VA.Pg4K.PDOffset;
  VAStart.Pg4K.PDOffset = VA.Pg4K.PDOffset;
  VAEnd.Pg4K.PDOffset = VA.Pg4K.PDOffset;
  //DBG ("PDE[%03x] at %p = %lx Region: %lx - %lx\n", VA.Pg4K.PDOffset, PDE, PDE->Uint64, VAStart.Uint64, VAEnd.Uint64);

  if (!PDE->Bits.Present) {
    //DBG ("-> Mapping not present!\n");
    return EFI_NO_MAPPING;
  }
  if (PDE->Bits.MustBeZero & 0x1) {
    // 2MB PDE
    PTE2M = (PAGE_TABLE_2M_ENTRY *)PDE;
    //DBG ("-> Nx:%x|G:%x|PAT:%x|D:%x|A:%x|PCD:%x|PWT:%x|US:%x|RW:%x|P:%x\n",
    //PTE2M->Bits.Nx, PTE2M->Bits.Global, PTE2M->Bits.PAT,
    //PTE2M->Bits.Dirty, PTE2M->Bits.Accessed,
    //PTE2M->Bits.CacheDisabled, PTE2M->Bits.WriteThrough,
    //PTE2M->Bits.UserSupervisor, PTE2M->Bits.ReadWrite, PTE2M->Bits.Present
    //);
    Start = (PTE2M->Uint64 & PT_ADDR_MASK_2M);
    *PhysicalAddr = Start + VA.Pg2M.PhysPgOffset;
    //DBG ("-> 2MB page %lx - %lx => %lx\n", Start, Start + 0x200000 - 1, *PhysicalAddr);
    return EFI_SUCCESS;
  }
  //DBG ("-> Nx:%x|A:%x|PCD:%x|PWT:%x|US:%x|RW:%x|P:%x -> %lx\n",
  //PDE->Bits.Nx, PDE->Bits.Accessed,
  //PDE->Bits.CacheDisabled, PDE->Bits.WriteThrough,
  //PDE->Bits.UserSupervisor, PDE->Bits.ReadWrite, PDE->Bits.Present,
  //(PDE->Uint64 & PT_ADDR_MASK_4K)
  //);

  // PTE
  PTE4K = (PAGE_TABLE_4K_ENTRY *)(PDE->Uint64 & PT_ADDR_MASK_4K);
  PTE4K += VA.Pg4K.PTOffset;
  VAStart.Pg4K.PTOffset = VA.Pg4K.PTOffset;
  VAEnd.Pg4K.PTOffset = VA.Pg4K.PTOffset;
  //DBG ("PTE[%03x] at %p = %lx Region: %lx - %lx\n", VA.Pg4K.PTOffset, PTE4K, PTE4K->Uint64, VAStart.Uint64, VAEnd.Uint64);

  if (!PTE4K->Bits.Present) {
    //DBG ("-> Mapping not present!\n");
    return EFI_NO_MAPPING;
  }
  //DBG ("-> Nx:%x|G:%x|PAT:%x|D:%x|A:%x|PCD:%x|PWT:%x|US:%x|RW:%x|P:%x -> %lx\n",
  //PTE4K->Bits.Nx, PTE4K->Bits.Global, PTE4K->Bits.PAT,
  //PTE4K->Bits.Dirty, PTE4K->Bits.Accessed,
  //PTE4K->Bits.CacheDisabled, PTE4K->Bits.WriteThrough,
  //PTE4K->Bits.UserSupervisor, PTE4K->Bits.ReadWrite, PTE4K->Bits.Present,
  //(PTE4K->Uint64 & PT_ADDR_MASK_4K)
  //);
  Start = (PTE4K->Uint64 & PT_ADDR_MASK_4K);
  *PhysicalAddr = Start + VA.Pg4K.PhysPgOffset;

  return EFI_SUCCESS;
}

/** Inits vm memory pool. Should be called while boot services are still usable. */
EFI_STATUS
VmAllocateMemoryPool () {
  EFI_STATUS            Status;
  EFI_PHYSICAL_ADDRESS  Addr;

  if (VmMemoryPool != NULL) {
    // already allocated
    return EFI_SUCCESS;
  }

  VmMemoryPoolFreePages = 0x200; // 2 MB should be enough
  Addr = 0x100000000; // max address

  Status = AllocatePagesFromTop (EfiBootServicesData, VmMemoryPoolFreePages, &Addr);
  if (Status != EFI_SUCCESS) {
    Print (L"VmAllocateMemoryPool: AllocatePagesFromTop (EfiBootServicesData) = %r\n", Status);
  } else {
    VmMemoryPool = (UINT8 *)Addr;
    //DBG ("VmMemoryPool = %lx - %lx\n", VmMemoryPool, VmMemoryPool + EFI_PAGES_TO_SIZE (VmMemoryPoolFreePages) - 1);
    //DBGnvr ("VmMemoryPool = %lx - %lx\n", VmMemoryPool, VmMemoryPool + EFI_PAGES_TO_SIZE (VmMemoryPoolFreePages) - 1);
  }

  return Status;
}

/** Central method for allocating pages for VM page maps. */
VOID *
VmAllocatePages (
  UINTN   NumPages
) {
  VOID    *AllocatedPages = NULL;
  /*
  EFI_STATUS        Status;
  EFI_PHYSICAL_ADDRESS  Addr;

  // for now, just allocate EfiBootServicesData as usuall.
  // if needed, we'll allocate somwhere from the top of mem.
  Status = gBS->AllocatePages (AllocateAnyPages, EfiBootServicesData, NumPages, &Addr);
  if (EFI_ERROR (Status)) {
    Addr = 0;
  }
  return (VOID *)Addr;
  */
  if (VmMemoryPoolFreePages >= (INTN)NumPages) {
    AllocatedPages = VmMemoryPool;
    VmMemoryPool += EFI_PAGES_TO_SIZE (NumPages);
    VmMemoryPoolFreePages -= NumPages;
  } else {
    //DBGnvr ("VmAllocatePages - no more pages!\n");
    CpuDeadLoop ();
  }

  return AllocatedPages;
}

/** Maps (remaps) 4K page given by VirtualAddr to PhysicalAddr page in PageTable. */
EFI_STATUS
VmMapVirtualPage (
  PAGE_MAP_AND_DIRECTORY_POINTER *PageTable,
  EFI_VIRTUAL_ADDRESS VirtualAddr,
  EFI_PHYSICAL_ADDRESS PhysicalAddr
) {
  EFI_PHYSICAL_ADDRESS            Start;
  VIRTUAL_ADDR                    VA, VAStart, VAEnd;
  PAGE_MAP_AND_DIRECTORY_POINTER  *PML4, *PDPE, *PDE;
  PAGE_TABLE_4K_ENTRY             *PTE4K;
  PAGE_TABLE_4K_ENTRY             *PTE4KTmp;
  PAGE_TABLE_2M_ENTRY             *PTE2M;
  PAGE_TABLE_1G_ENTRY             *PTE1G;
  UINTN                           Index;

  VA.Uint64 = (UINT64)VirtualAddr;
  //VA_FIX_SIGN_EXTEND (VA);
  //DBG ("VmMapVirtualPage VA %lx => PA %lx\nPageTable: %p\n", VirtualAddr, PhysicalAddr, PageTable);
  //DBG ("VA: %lx => Indexes PML4=%x, PDP=%x, PD=%x, PT=%x\n",
  //  VA.Uint64, VA.Pg4K.PML4Offset, VA.Pg4K.PDPOffset, VA.Pg4K.PDOffset, VA.Pg4K.PTOffset);

  // PML4
  PML4 = PageTable;
  PML4 += VA.Pg4K.PML4Offset;
  // there is a problem if our PML4 points to the same table as first PML4 entry
  // since we may mess the mapping of first virtual region (happens in VBox and probably DUET).
  // check for this on first call and if true, just clear our PML4 - we'll rebuild it in later step
  if (PML4 != PageTable && PML4->Bits.Present && PageTable->Bits.PageTableBaseAddress == PML4->Bits.PageTableBaseAddress) {
    //DBG ("PML4 points to the same table as first PML4 - releasing it and rebuiding in a separate table\n");
    PML4->Uint64 = 0;
  }

  // prepare region start and end
  VAStart.Uint64 = 0;
  VAStart.Pg4K.PML4Offset = VA.Pg4K.PML4Offset;
  VA_FIX_SIGN_EXTEND (VAStart);
  VAEnd.Uint64 = ~(UINT64)0;
  VAEnd.Pg4K.PML4Offset = VA.Pg4K.PML4Offset;
  VA_FIX_SIGN_EXTEND (VAEnd);
  // print it
  //DBG ("PML4[%03x] at %p = %lx Region: %lx - %lx\n", VA.Pg4K.PML4Offset, PML4, PML4->Uint64, VAStart.Uint64, VAEnd.Uint64);

  if (!PML4->Bits.Present) {
    //DBG ("-> Mapping not present, creating new PML4 entry and page with PDPE entries!\n");
    PDPE = (PAGE_MAP_AND_DIRECTORY_POINTER *)VmAllocatePages (1);
    if (PDPE == NULL) {
      //DBG ("No memory - exiting.\n");
      return EFI_NO_MAPPING;
    }

    ZeroMem (PDPE, EFI_PAGE_SIZE);
    // init this whole 512GB region with 512 1GB entry pages to map first 512GB phys space
    PTE1G = (PAGE_TABLE_1G_ENTRY *)PDPE;
    Start = 0;
    for (Index = 0; Index < 512; Index++) {
      PTE1G->Uint64 = Start & PT_ADDR_MASK_1G;
      PTE1G->Bits.ReadWrite = 1;
      PTE1G->Bits.Present = 1;
      PTE1G->Bits.MustBe1 = 1;
      PTE1G++;
      Start += 0x40000000;
    }

    // put it to PML4
    PML4->Uint64 = ((UINT64)PDPE) & PT_ADDR_MASK_4K;
    PML4->Bits.ReadWrite = 1;
    PML4->Bits.Present = 1;
    //DBG ("added to PLM4 as %lx\n", PML4->Uint64);
    // and continue with mapping ...
  }
  //DBG ("-> Nx:%x|A:%x|PCD:%x|PWT:%x|US:%x|RW:%x|P:%x -> %lx\n",
  //PML4->Bits.Nx, PML4->Bits.Accessed,
  //PML4->Bits.CacheDisabled, PML4->Bits.WriteThrough,
  //PML4->Bits.UserSupervisor, PML4->Bits.ReadWrite, PML4->Bits.Present,
  //(PML4->Uint64 & PT_ADDR_MASK_4K)
  //);

  // PDPE
  PDPE = (PAGE_MAP_AND_DIRECTORY_POINTER *)(PML4->Uint64 & PT_ADDR_MASK_4K);
  PDPE += VA.Pg4K.PDPOffset;
  VAStart.Pg4K.PDPOffset = VA.Pg4K.PDPOffset;
  VAEnd.Pg4K.PDPOffset = VA.Pg4K.PDPOffset;
  //DBG ("PDPE[%03x] at %p = %lx Region: %lx - %lx\n", VA.Pg4K.PDPOffset, PDPE, PDPE->Uint64, VAStart.Uint64, VAEnd.Uint64);

  if (!PDPE->Bits.Present || (PDPE->Bits.MustBeZero & 0x1)) {
    //DBG ("-> Mapping not present or mapped as 1GB page, creating new PDPE entry and page with PDE entries!\n");
    PDE = (PAGE_MAP_AND_DIRECTORY_POINTER *)VmAllocatePages (1);
    if (PDE == NULL) {
      //DBG ("No memory - exiting.\n");
      return EFI_NO_MAPPING;
    }
    ZeroMem (PDE, EFI_PAGE_SIZE);

    if (PDPE->Bits.MustBeZero & 0x1) {
      // was 1GB page - init new PDE array to get the same mapping but with 2MB pages
      //DBG ("-> was 1GB page, initing new PDE array to get the same mapping but with 2MB pages!\n");
      PTE2M = (PAGE_TABLE_2M_ENTRY *)PDE;
      Start = (PDPE->Uint64 & PT_ADDR_MASK_1G);
      for (Index = 0; Index < 512; Index++) {
        PTE2M->Uint64 = Start & PT_ADDR_MASK_2M;
        PTE2M->Bits.ReadWrite = 1;
        PTE2M->Bits.Present = 1;
        PTE2M->Bits.MustBe1 = 1;
        PTE2M++;
        Start += 0x200000;
      }
    }

    // put it to PDPE
    PDPE->Uint64 = ((UINT64)PDE) & PT_ADDR_MASK_4K;
    PDPE->Bits.ReadWrite = 1;
    PDPE->Bits.Present = 1;
    //DBG ("added to PDPE as %lx\n", PDPE->Uint64);
    // and continue with mapping ...
  }

  //DBG ("-> Nx:%x|A:%x|PCD:%x|PWT:%x|US:%x|RW:%x|P:%x -> %lx\n",
  //PDPE->Bits.Nx, PDPE->Bits.Accessed,
  //PDPE->Bits.CacheDisabled, PDPE->Bits.WriteThrough,
  //PDPE->Bits.UserSupervisor, PDPE->Bits.ReadWrite, PDPE->Bits.Present,
  //(PDPE->Uint64 & PT_ADDR_MASK_4K)
  //);

  // PDE
  PDE = (PAGE_MAP_AND_DIRECTORY_POINTER *)(PDPE->Uint64 & PT_ADDR_MASK_4K);
  PDE += VA.Pg4K.PDOffset;
  VAStart.Pg4K.PDOffset = VA.Pg4K.PDOffset;
  VAEnd.Pg4K.PDOffset = VA.Pg4K.PDOffset;
  //DBG ("PDE[%03x] at %p = %lx Region: %lx - %lx\n", VA.Pg4K.PDOffset, PDE, PDE->Uint64, VAStart.Uint64, VAEnd.Uint64);

  if (!PDE->Bits.Present || (PDE->Bits.MustBeZero & 0x1)) {
    //DBG ("-> Mapping not present or mapped as 2MB page, creating new PDE entry and page with PTE4K entries!\n");
    PTE4K = (PAGE_TABLE_4K_ENTRY *)VmAllocatePages (1);
    if (PTE4K == NULL) {
      //DBG ("No memory - exiting.\n");
      return EFI_NO_MAPPING;
    }

    ZeroMem (PTE4K, EFI_PAGE_SIZE);

    if (PDE->Bits.MustBeZero & 0x1) {
      // was 2MB page - init new PTE array to get the same mapping but with 4KB pages
      //DBG ("-> was 2MB page - initing new PTE array to get the same mapping but with 4KB pages!\n");
      PTE4KTmp = (PAGE_TABLE_4K_ENTRY *)PTE4K;
      Start = (PDE->Uint64 & PT_ADDR_MASK_2M);
      for (Index = 0; Index < 512; Index++) {
        PTE4KTmp->Uint64 = Start & PT_ADDR_MASK_4K;
        PTE4KTmp->Bits.ReadWrite = 1;
        PTE4KTmp->Bits.Present = 1;
        PTE4KTmp++;
        Start += 0x1000;
      }
    }

    // put it to PDE
    PDE->Uint64 = ((UINT64)PTE4K) & PT_ADDR_MASK_4K;
    PDE->Bits.ReadWrite = 1;
    PDE->Bits.Present = 1;
    //DBG ("added to PDE as %lx\n", PDE->Uint64);
    // and continue with mapping ...
  }

  //DBG ("-> Nx:%x|A:%x|PCD:%x|PWT:%x|US:%x|RW:%x|P:%x -> %lx\n",
  //PDE->Bits.Nx, PDE->Bits.Accessed,
  //PDE->Bits.CacheDisabled, PDE->Bits.WriteThrough,
  //PDE->Bits.UserSupervisor, PDE->Bits.ReadWrite, PDE->Bits.Present,
  //(PDE->Uint64 & PT_ADDR_MASK_4K)
  //);

  // PTE
  PTE4K = (PAGE_TABLE_4K_ENTRY *)(PDE->Uint64 & PT_ADDR_MASK_4K);
  PTE4K += VA.Pg4K.PTOffset;
  VAStart.Pg4K.PTOffset = VA.Pg4K.PTOffset;
  VAEnd.Pg4K.PTOffset = VA.Pg4K.PTOffset;
  //DBG ("PTE[%03x] at %p = %lx Region: %lx - %lx\n", VA.Pg4K.PTOffset, PTE4K, PTE4K->Uint64, VAStart.Uint64, VAEnd.Uint64);
  if (PTE4K->Bits.Present) {
    //DBG ("mapping already present - remapping!\n");
  }

  //DBG ("-> Nx:%x|G:%x|PAT:%x|D:%x|A:%x|PCD:%x|PWT:%x|US:%x|RW:%x|P:%x -> %lx\n",
   //PTE4K->Bits.Nx, PTE4K->Bits.Global, PTE4K->Bits.PAT,
   //PTE4K->Bits.Dirty, PTE4K->Bits.Accessed,
   //PTE4K->Bits.CacheDisabled, PTE4K->Bits.WriteThrough,
   //PTE4K->Bits.UserSupervisor, PTE4K->Bits.ReadWrite, PTE4K->Bits.Present,
   //(PTE4K->Uint64 & PT_ADDR_MASK_4K)
   //);
  // put it to PTE
  PTE4K->Uint64 = ((UINT64)PhysicalAddr) & PT_ADDR_MASK_4K;
  PTE4K->Bits.ReadWrite = 1;
  PTE4K->Bits.Present = 1;
  //DBG ("added to PTE4K as %lx\n", PTE4K->Uint64);

  return EFI_SUCCESS;
}

/** Maps (remaps) NumPages 4K pages given by VirtualAddr to PhysicalAddr pages in PageTable. */
EFI_STATUS
VmMapVirtualPages (
  PAGE_MAP_AND_DIRECTORY_POINTER  *PageTable,
  EFI_VIRTUAL_ADDRESS             VirtualAddr,
  UINTN                           NumPages,
  EFI_PHYSICAL_ADDRESS            PhysicalAddr
) {
  EFI_STATUS    Status = EFI_SUCCESS;

  while ((NumPages > 0) && (Status == EFI_SUCCESS)) {
    Status = VmMapVirtualPage (PageTable, VirtualAddr, PhysicalAddr);
    VirtualAddr += 0x1000;
    PhysicalAddr += 0x1000;
    NumPages--;
    //DBG ("NumPages: %d, %lx => %lx\n", NumPages, VirtualAddr, PhysicalAddr);
    //WaitForKeyPress (L"press a key to continue\n");
  }

  return Status;
}

/** Flashes TLB caches. */
VOID
VmFlashCaches () {
  // just reload CR3
  AsmWriteCr3 (AsmReadCr3 ());
}
